from pointnet_util import index_points, square_distance
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class TransformerBlock(nn.Module):
    def __init__(self, d_points, d_model, k) -> None:
        super().__init__()
        self.fc1 = nn.Linear(d_points, d_model)
        self.fc2 = nn.Linear(d_model, d_points)
        self.fc_delta = nn.Sequential(
            nn.Linear(3, d_model),
            nn.ReLU(),
            nn.Linear(d_model, d_model)
        )
        self.fc_gamma = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.ReLU(),
            nn.Linear(d_model, d_model)
        )
        self.w_qs = nn.Linear(d_model, d_model, bias=False)
        self.w_ks = nn.Linear(d_model, d_model, bias=False)
        self.w_vs = nn.Linear(d_model, d_model, bias=False)
        self.k = k

    # xyz: b x n x 3, features: b x n x f
    def forward(self, xyz, features):
        dists = square_distance(xyz, xyz)  # [B, N, N]
        knn_idx = dists.argsort()[:, :, :self.k]  # b x n x k [B, N, k]选取最近的k个点的地址
        knn_xyz = index_points(xyz, knn_idx)  # [b n k 3](已经验证)
        
        pre = features
        x = self.fc1(features)
        q, k, v = self.w_qs(x), index_points(self.w_ks(x), knn_idx), index_points(self.w_vs(x), knn_idx)
        '''
        index_points
        Input:
            points: input points data, [B, N, C]
            idx: sample index data, [B, S, [K]]
        Return:
            new_points:, indexed points data, [B, S, [K], C]
        index_points函数是按照输入的点云数据和索引返回由索引的点云数据。
        例如points为B×2048×3的点云，idx为[1,333,1000,2000](S的维度)，
        则返回B个样本中每个样本的第1,333,1000,2000个点组成的B×4×3的点云集。
        当然如果idx为一个[B,S]维度的，则它会按照idx中的维度结构将其提取成[B,S,C]。
        '''

        pos_enc = self.fc_delta(xyz[:, :, None] - knn_xyz)
        # xyz[:, :, None]增加1维，变为：[B,N,1,3]，广播后相减 输出b x n x k x f
        attn = self.fc_gamma(q[:, :, None] - k + pos_enc)
        attn = F.softmax(attn / np.sqrt(k.size(-1)), dim=-2)  # b x n x k x f
        #softmax(q-k)/(d^1/2),再与v做爱因斯坦积
        res = torch.einsum('bmnf,bmnf->bmf', attn, v + pos_enc)  #[B, N, d]
        res = self.fc2(res) + pre
        return res, attn  # [B, N, D] [B N K D]

if __name__ == '__main__':
    import time
    #device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

    d_in = 7
    d_model = 23
    # cloud1 = 1000 * torch.randn(1, 447, d_in).to(device)
    cloud1 = 1000 * torch.randn(1, 447, d_in) #
    cloud2 = cloud1[:,:,:3]
    model = TransformerBlock(d_in, d_model, 4)
    # model.load_state_dict(torch.load('checkpoints/checkpoint_100.pth'))
    model.eval()

    t0 = time.time()
    with torch.no_grad():
        pred = model(cloud2,cloud1)
        print(pred)
    t1 = time.time()
    # print(pred)
    print(t1 - t0)